package com.troila.xfspark.websocket;

import com.troila.xfspark.model.SparkResponse;
import org.springframework.web.socket.WebSocketSession;

import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

public class WebSocketSessionManager {

    private static final ConcurrentHashMap<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    private static final ConcurrentHashMap<String, List<SparkResponse.SparkChoices>> choicesMap = new ConcurrentHashMap<>();

    private static WebSocketSession sparkSession = null;

    private WebSocketSessionManager() {
    }

    public static WebSocketSession put(String sessionId, WebSocketSession webSocketSession) {
        return sessions.put(sessionId, webSocketSession);
    }

    public static List<WebSocketSession> findAllSessions() {
        return sessions.entrySet().stream()
                .map(entry -> entry.getValue())
                .collect(Collectors.toList());
    }

    public static Set<String> findAllSessionIds() {
        return sessions.keySet();
    }

    public static WebSocketSession get(String sessionId) {
        return sessions.get(sessionId);
    }

    public static WebSocketSession remove(String sessionId) {
        return sessions.remove(sessionId);
    }

    public static WebSocketSession getSparkSession() {
        return sparkSession;
    }

    public static void setSparkSession(WebSocketSession sparkSession) {
        WebSocketSessionManager.sparkSession = sparkSession;
    }

    public static void putChoices(String sessionId, SparkResponse.SparkChoices choice) {
        if (choicesMap.containsKey(sessionId)) {
            choicesMap.get(sessionId).add(choice);
        } else {
            List<SparkResponse.SparkChoices> choiceList = new ArrayList<>();
            choiceList.add(choice);
            choicesMap.put(sessionId, choiceList);
        }
    }

    public static List<SparkResponse.SparkChoices> getChoices(String sessionId) {
        return choicesMap.get(sessionId);
    }

    public static void removeChoices(String sessionId) {
        choicesMap.remove(sessionId);
    }
}
